import sys
import random
import time
from PyQt5 import QtWidgets
from PyQt5.QtWidgets import QFileDialog, QMessageBox, QListWidget, QDialog
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from src import solver as Solver
from src import visualizer as Ver
import pandas as pd

class RiemannGUI(QtWidgets.QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("黎曼问题求解器")
        self.setGeometry(100, 100, 1600, 1200)  # 调整窗口大小

        central_widget = QtWidgets.QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QtWidgets.QVBoxLayout(central_widget)

        # ====================== 参数输入区域 ======================
        param_frame = QtWidgets.QGroupBox("参数输入")
        param_layout = QtWidgets.QHBoxLayout(param_frame)

        # 左侧：分片数量选择和参数输入
        left_layout = QtWidgets.QVBoxLayout()
        self.tab_widget = QtWidgets.QTabWidget()
        self.two_region_tab = QtWidgets.QWidget()
        self.three_region_tab = QtWidgets.QWidget()
        self.tab_widget.addTab(self.two_region_tab, "两段常数")
        self.tab_widget.addTab(self.three_region_tab, "三段常数")
        left_layout.addWidget(self.tab_widget)

        # 两段常数参数界面
        self.setup_two_region_ui()
        # 三段常数参数界面
        self.setup_three_region_ui()

        # 右侧：文件导入功能
        right_layout = QtWidgets.QVBoxLayout()
        self.import_btn = QtWidgets.QPushButton("导入网格点文件")
        self.import_btn.clicked.connect(self.import_points_file)
        self.clear_btn = QtWidgets.QPushButton("清空数据")
        self.clear_btn.clicked.connect(self.clear_data)
        right_layout.addWidget(QtWidgets.QLabel("导入网格点(不导入视为自适应生成网格)"))
        right_layout.addWidget(self.import_btn)
        right_layout.addWidget(self.clear_btn)

        param_layout.addLayout(left_layout)
        param_layout.addLayout(right_layout)
        main_layout.addWidget(param_frame)

        # ====================== 功能按钮区域 ======================
        btn_frame = QtWidgets.QGroupBox("操作")
        btn_layout = QtWidgets.QHBoxLayout(btn_frame)
        self.calculate_btn = QtWidgets.QPushButton("计算")
        self.calculate_btn.clicked.connect(self.calculate)
        self.memory_btn = QtWidgets.QPushButton("记忆")
        self.memory_btn.clicked.connect(self.show_memory)
        self.plot_btn = QtWidgets.QPushButton("绘制求解结果")
        self.plot_btn.clicked.connect(self.plot_graph)
        self.export_btn = QtWidgets.QPushButton("导出为csv")
        self.export_btn.clicked.connect(self.export_to_csv)
        btn_layout.addWidget(self.calculate_btn)
        btn_layout.addWidget(self.memory_btn)
        btn_layout.addWidget(self.plot_btn)
        btn_layout.addWidget(self.export_btn)
        main_layout.addWidget(btn_frame)

        # ====================== 日志框 ======================
        self.log_text = QtWidgets.QTextEdit()
        self.log_text.setReadOnly(True)
        main_layout.addWidget(QtWidgets.QLabel("日志信息"))
        main_layout.addWidget(self.log_text)

        # ====================== 图形显示区域 ======================
        self.figure = plt.figure(figsize=(12, 5))
        self.canvas = FigureCanvas(self.figure)
        self.toolbar = NavigationToolbar(self.canvas, self)
        main_layout.addWidget(self.toolbar)
        main_layout.addWidget(self.canvas)

        # 初始化变量
        self.current_points = None  # 当前使用的网格点数据
        self.set_random_initial_values()
        self.current_sol = None  # 当前解析解
        self.current_cal = None  # 当前数值解
        self.memory_list = []  # 记忆列表
        self.is_calculating = False  # 计算状态标志
        self.is_plotting = False  # 绘图状态标志
        self.import_btn.setEnabled(True)  # 初始化导入按钮状态

    def setup_two_region_ui(self):
        """两段常数参数界面"""
        layout = QtWidgets.QFormLayout(self.two_region_tab)
        self.x0_input = QtWidgets.QComboBox()
        self.ul__input = QtWidgets.QComboBox()
        self.ur__input = QtWidgets.QComboBox()
        for i in range(-5, 6):
            self.x0_input.addItem(str(i), i)
            self.ul__input.addItem(str(i), i)
            self.ur__input.addItem(str(i), i)
        layout.addRow("x₀:", self.x0_input)
        layout.addRow("uₗ:", self.ul__input)
        layout.addRow("uᵣ:", self.ur__input)

    def setup_three_region_ui(self):
        """三段常数参数界面"""
        layout = QtWidgets.QFormLayout(self.three_region_tab)
        self.xl_input = QtWidgets.QComboBox()
        self.xr_input = QtWidgets.QComboBox()
        self.ul_input = QtWidgets.QComboBox()
        self.um_input = QtWidgets.QComboBox()
        self.ur_input = QtWidgets.QComboBox()
        for i in range(-5, 6):
            self.xl_input.addItem(str(i), i)
            self.xr_input.addItem(str(i), i)
            self.ul_input.addItem(str(i), i)
            self.um_input.addItem(str(i), i)
            self.ur_input.addItem(str(i), i)
        layout.addRow("xₗ:", self.xl_input)
        layout.addRow("xᵣ:", self.xr_input)
        layout.addRow("uₗ:", self.ul_input)
        layout.addRow("uₘ:", self.um_input)
        layout.addRow("uᵣ:", self.ur_input)

    def set_random_initial_values(self):
        """设置随机初值"""
        def random_idx(): return random.randint(0, 10)
        # 两段参数
        self.x0_input.setCurrentIndex(random_idx())
        self.ul__input.setCurrentIndex(random_idx())
        self.ur__input.setCurrentIndex(random_idx())
        # 三段参数
        self.xl_input.setCurrentIndex(random_idx())
        self.xr_input.setCurrentIndex(random_idx())
        self.ul_input.setCurrentIndex(random_idx())
        self.um_input.setCurrentIndex(random_idx())
        self.ur_input.setCurrentIndex(random_idx())

    def import_points_file(self):
        """导入自定义网格点文件（带按钮禁用和日志时间）"""
        filename, _ = QFileDialog.getOpenFileName(self, "导入文件", "", "Excel Files (*.xlsx)")
        if filename:
            try:
                # 禁用导入按钮防止重复点击
                self.import_btn.setEnabled(False)
                start_time = time.time()
                
                # 记录开始日志
                self.log_text.append(f"开始导入网格点文件：{filename}")
                
                # 执行导入和排序
                self.current_points = pd.read_excel(filename)
                self.current_points.sort_values(by=['x', 't'], inplace=True)
                self.current_points.reset_index(drop=True, inplace=True)
                
                # 计算耗时并记录结束日志
                end_time = time.time()
                elapsed_time = end_time - start_time
                self.log_text.append(f"网格点导入并排序完成，耗时：{elapsed_time:.4f} 秒")
                QMessageBox.information(self, "成功", "网格点文件导入并排序完成")
                
            except Exception as e:
                QMessageBox.warning(self, "错误", f"文件导入失败：{str(e)}")
            finally:
                # 恢复按钮状态
                self.import_btn.setEnabled(True)

    def clear_data(self):
        """清空当前数据"""
        self.current_points = None  # 恢复默认网格点
        QMessageBox.information(self, "提示", "数据已恢复为默认网格点")

    def calculate(self):
        """计算函数"""
        if self.is_calculating:
            return
        self.is_calculating = True
        self.calculate_btn.setEnabled(False)  # 禁用计算按钮

        try:
            # 获取当前参数
            if self.tab_widget.currentIndex() == 0:
                # 两段常数情况
                x0 = self.x0_input.currentData()
                ul = self.ul__input.currentData()
                ur = self.ur__input.currentData()
                params = (x0, ul, ur)
                self.log_text.append(f"当前参数（两段常数）：x₀={x0}, uₗ={ul}, uᵣ={ur}")
            else:
                # 三段常数情况
                xl = self.xl_input.currentData()
                xr = self.xr_input.currentData()
                ul = self.ul_input.currentData()
                um = self.um_input.currentData()
                ur = self.ur_input.currentData()
                params = (xl, xr, ul, um, ur)
                self.log_text.append(f"当前参数（三段常数）：xₗ={xl}, xᵣ={xr}, uₗ={ul}, uₘ={um}, uᵣ={ur}")

            self.log_text.append("开始计算...")
            start_time = time.time()

            # 求解解析解
            if self.tab_widget.currentIndex() == 0:
                self.current_sol = Solver.solve_2(x0, ul, ur)
            else:
                self.current_sol = Solver.solve_3(xl, xr, ul, um, ur)

            # 计算数值解
            if self.current_points is not None:
                self.current_cal = Solver.calculate(self.current_sol, self.current_points)
            else:
                self.current_cal=None
                self.log_text.append("无坐标网格，跳过部分计算")

            end_time = time.time()
            elapsed_time = end_time - start_time
            self.log_text.append(f"计算结束，用时：{elapsed_time:.4f} 秒")
            self.log_text.append(f"当前解参数：{params}")

            # 将当前解加入记忆列表
            self.memory_list.append((params, self.current_sol, self.current_cal))

        except Exception as e:
            QMessageBox.warning(self, "计算失败", f"计算错误：{str(e)}")
        finally:
            self.is_calculating = False
            self.calculate_btn.setEnabled(True)  # 启用计算按钮

    def show_memory(self):
        """显示记忆列表窗口"""
        dialog = QDialog(self)
        dialog.setWindowTitle("记忆列表")
        layout = QtWidgets.QVBoxLayout(dialog)
        list_widget = QListWidget()
        for params, _, _ in self.memory_list:
            list_widget.addItem(str(params))
        layout.addWidget(list_widget)
        select_btn = QtWidgets.QPushButton("选择")
        select_btn.clicked.connect(lambda: self.select_memory(list_widget.currentRow()))
        layout.addWidget(select_btn)
        dialog.exec_()

    def select_memory(self, index):
        """选择记忆中的解"""
        if index >= 0 and index < len(self.memory_list):
            params, self.current_sol, self.current_cal = self.memory_list[index]
            self.log_text.append(f"当前解参数：{params}")

    def plot_graph(self):
        """绘制图形"""
        if self.is_plotting:
            return
        if self.current_sol is None:
            QMessageBox.warning(self, "绘图失败", "请先进行计算")
            return
        self.is_plotting = True
        self.plot_btn.setEnabled(False)  # 禁用绘图按钮
        try:
            self.log_text.append("正在绘图...")
            start_time = time.time()

            # 调用可视化函数生成新的Figure
            new_figure = Ver.show(self.current_sol, self.current_cal, False)

            # 替换当前画布为新生成的Figure
            # 移除旧画布和工具栏
            old_canvas = self.canvas
            old_toolbar = self.toolbar
            layout = self.centralWidget().layout()
            layout.removeWidget(old_canvas)
            layout.removeWidget(old_toolbar)
            old_canvas.deleteLater()
            old_toolbar.deleteLater()

            # 创建新画布和工具栏并关联到新Figure
            self.canvas = FigureCanvas(new_figure)
            self.toolbar = NavigationToolbar(self.canvas, self)
            layout.addWidget(self.toolbar)
            layout.addWidget(self.canvas)

            # 刷新显示
            self.canvas.draw()

            end_time = time.time()
            elapsed_time = end_time - start_time
            self.log_text.append(f"绘图结束，用时：{elapsed_time:.4f} 秒")

        except Exception as e:
            QMessageBox.warning(self, "绘图失败", f"计算错误：{str(e)}")
        finally:
            self.is_plotting = False
            self.plot_btn.setEnabled(True)  # 启用绘图按钮

    def export_to_csv(self):
        """导出结果"""
        if self.current_cal is None:
            QMessageBox.warning(self, "导出失败", "请先进行计算或导入坐标网格")
            return
        try:
            params = self.current_sol.argv
            filename = f"data/result/result_{'_'.join(map(str, params))}.csv"
            self.current_cal.to_csv(filename, index=False)
            QMessageBox.information(self, "成功", f"结果已导出至：{filename}")
        except Exception as e:
            QMessageBox.warning(self, "导出失败", f"文件保存错误：{str(e)}")

if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    gui = RiemannGUI()
    gui.show()
    sys.exit(app.exec_())